"""
This module contains functions related to hydropower technologies.
"""
import pyoptinterface as poi
from loguru import logger

class AddHydroConstraints:
    """Class for hydropower constraints and calculations.
    """
    def __init__(self, model : object) -> None:
        self.model = model
        # Define variables and constraints
        model.inflow_cons = poi.make_tupledict(
            model.reservoir, model.time, 
            rule=self.inflow_rule
        )
        model.outflow_cons = poi.make_tupledict(
            model.reservoir, model.time, 
            rule=self.outflow_rule
        )
        model.supply_demand_cons = poi.make_tupledict(
            model.water_user, model.time, 
            rule=self.supply_demand_rule
        )
        model.init_storage_cons = poi.make_tupledict(
            model.reservoir, 
            rule=self.init_storage_rule
        )
        model.storage_low_cons = poi.make_tupledict(
            model.reservoir, model.time_p, 
            rule=self.storage_low_bound_rule
        )
        model.storage_up_cons = poi.make_tupledict(
            model.reservoir, model.time_p, 
            rule=self.storage_up_bound_rule
        )
        model.water_balance_cons = poi.make_tupledict(
            model.reservoir, model.time, 
            rule=self.water_balance_rule
        )
        model.spill_up_cons = poi.make_tupledict(
            model.spillflow_up_cons_set, model.time, 
            rule=self.spill_up_bound_rule
        )
        model.spill_down_cons = poi.make_tupledict(
            model.spillflow_down_cons_set, model.time, 
            rule=self.spill_down_bound_rule
        )

        # 引入/调出水量的上限约束（仅对已配置的(source, target, time)三元组施加）
        model.import_export_up_cons = poi.make_tupledict(
            model.import_export_up_cons_set,
            rule=self.import_export_up_rule
        )
        # 引入/调出水量的下限约束（仅对已配置的(source, target, time)三元组施加）
        model.import_export_down_cons = poi.make_tupledict(
            model.import_export_down_cons_set,
            rule=self.import_export_down_rule
        )
        
        # 汛末库容约束（仅对已配置的水库施加等于约束）
        model.end_storage_flood_season_cons = poi.make_tupledict(
            model.end_storage_flood_season_cons_set,
            rule=self.end_storage_flood_season_rule
        )
        
        # 调度期末库容约束（仅对已配置的水库施加等于约束）
        model.end_storage_dispatch_term_cons = poi.make_tupledict(
            model.end_storage_dispatch_term_cons_set,
            rule=self.end_storage_dispatch_term_rule
        )

    def inflow_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Total inflow of reservoir is equal to the sum of importflow and 
        supply.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time.

        Returns
        -------
        poi.ConstraintIndex
            Total inflow of reservoir.
        """
        model = self.model
        
        # 查找该水库对应的上游水库（入流加上上游水库的出流）
        up_streams = model.param['reservoir_connections'].get(r, [])
        up_streams_flow = poi.ExprBuilder()
        for ups in up_streams:
            up_streams_flow += model.spillflow[ups, t]

        # 查找该水库是否有引入水的关系（入流加上引水水的水量）
        import_reservoirs = model.param['import_relationship'].get(r, [])
        import_reservoirs_flow = poi.ExprBuilder()
        for ims in import_reservoirs:
            import_reservoirs_flow += model.importflow[ims, r, t]

        lhs = model.inflow[r, t] - (up_streams_flow + model.param['runoff'][r, t] + import_reservoirs_flow)
        return model.add_linear_constraint(lhs, poi.Eq, 0)
    
    def outflow_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Total outflow of reservoir is equal to the sum of spillage and 
        supply.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time.
            
        Returns
        -------
        poi.ConstraintIndex
            Total outflow of reservoir.
        """
        model = self.model
        # 查找这座水库是否会调出水（如大伙房 向 碧流河引出水，但仅针对集团内部水库互相引出水）
        sources = model.param['import_relationship_source'].get(r, [])
        source_flow = poi.ExprBuilder()
        for src in sources:
            source_flow += model.importflow[r, src, t]

        # 查找这座水库都需要给哪些用水户供水
        supplies = model.param['reservoir_supply'].get(r, [])
        supply_flow = poi.ExprBuilder()
        for usr in supplies:
            supply_flow += model.supplyflow[r, usr, t]

        lhs = model.outflow[r, t] - (model.spillflow[r, t] + supply_flow + source_flow)
        return model.add_linear_constraint(lhs, poi.Eq, 0)
    
    def supply_demand_rule(self, u : str, t : str) -> poi.ConstraintIndex:
        """Supply demand balance.

        Parameters
        ----------
        u : str
            User name.
        t : str
            Time.
            
        Returns
        -------
        poi.ConstraintIndex
            Supply demand balance.
        """
        model = self.model
        # dt = model.param['dt']
        lhs = poi.ExprBuilder()
        lhs += poi.quicksum(
            model.supplyflow[r, u, t] 
            for r in model.reservoir if model.param['supply_relationship'][r, u] == 1
        )
        lhs -= model.param['demand'][u, t] * 1e-6
        return model.add_linear_constraint(lhs, poi.Eq, 0)

    def init_storage_rule(self, r : str) -> poi.ConstraintIndex:
        """Determine storage of reservoir in the initial time.

        Parameters
        ----------
        r : str
            Reservoir name.
            
        Returns
        -------
        poi.ConstraintIndex
            The constraint of the model.
        """
        model = self.model
        lhs = model.storage[r, model.time[0]] - model.param['storage_init'][r, model.time[0]]
        return model.add_linear_constraint(lhs, poi.Eq, 0)

    def end_storage_flood_season_rule(self, r: str) -> poi.ConstraintIndex:
        """End storage constraint at flood season end (October).

        Parameters
        ----------
        r : str
            Reservoir name.

        Returns
        -------
        poi.ConstraintIndex
            Equality constraint for storage at flood season end.
        """
        model = self.model
        # 获取汛末时间（10月份）
        october_times = [t for t in model.time if isinstance(t, str) and t.endswith('-10')]
        if not october_times:
            raise ValueError(f"汛末库容约束要求有10月份的数据，但在时间序列中未找到")
        # 选择最新的一个10月（如果跨年，则取年份最大的10月）
        flood_end_time = sorted(october_times)[-1]
        
        # 等于约束: storage[r, flood_end_time] == target_storage
        target_val = model.param['end_storage_flood_season_cons'][r]
        lhs = model.storage[r, flood_end_time] - target_val
        return model.add_linear_constraint(lhs, poi.Eq, 0)

    def end_storage_dispatch_term_rule(self, r: str) -> poi.ConstraintIndex:
        """End storage constraint at dispatch term end.

        Parameters
        ----------
        r : str
            Reservoir name.

        Returns
        -------
        poi.ConstraintIndex
            Equality constraint for storage at dispatch term end.
        """
        model = self.model
        # 获取调度期末时间（最后一个时间步）
        dispatch_end_time = model.time_p[-1]
        
        # 等于约束: storage[r, dispatch_end_time] == target_storage
        target_val = model.param['end_storage_dispatch_term_cons'][r]
        lhs = model.storage[r, dispatch_end_time] - target_val
        return model.add_linear_constraint(lhs, poi.Eq, 0)

    def storage_low_bound_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Lower bound of storage of reservoir.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time period.
            
        Returns
        -------
        poi.ConstraintIndex
            Lower bound of storage of reservoir.
        """
        model = self.model
        lhs = model.storage[r, t] - model.param['storage_down'][r, t]
        return model.add_linear_constraint(lhs, poi.Geq, 0)
        
    def storage_up_bound_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Upper bound of storage of reservoir.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time period.
            
        Returns
        -------
        poi.ConstraintIndex
            Upper bound of storage of reservoir.
        """
        model = self.model
        lhs = model.storage[r, t] - model.param['storage_up'][r, t]
        return model.add_linear_constraint(lhs, poi.Leq, 0)

    def water_balance_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Water balance of reservoir.

        Parameters
        ----------
        r : str
            Reservoir name.
        t : str
            Time.
            
        Returns
        -------
        poi.ConstraintIndex
            Water balance of reservoir.
        """
        model = self.model
        lhs = poi.ExprBuilder()

        # 找到当前时间在time_p中的索引
        time_index = model.time_p.index(t)
        next_time = model.time_p[time_index + 1]  # 下一个时间点
    
        # 初始库容
        lhs += model.storage[r, t]
        # 净流入量（均为水量）
        lhs += (model.inflow[r, t] - model.outflow[r, t])
        # 末库容 = 初始库容 + 净流入量 - 弃水量
        lhs -= model.storage[r, next_time]
        return model.add_linear_constraint(lhs, poi.Eq, 0, name=f'water_balance_{r}_{t}')

    def spill_up_bound_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Upper bound of spillflow of reservoir.

        Parameters
        ----------
        r : str
            Set spillflow_up Reservoir name.
        t : str
            Time period.
            
        Returns
        -------
        poi.ConstraintIndex
            Upper bound of spillflow of reservoir.
        """
        model = self.model
        lhs = model.spillflow[r, t] - model.param['spillflow_up'][r, t]
        return model.add_linear_constraint(lhs, poi.Leq, 0)

    def spill_down_bound_rule(self, r : str, t : str) -> poi.ConstraintIndex:
        """Lower bound of spillflow of reservoir.

        Parameters
        ----------
        r : str
            Set spillflow_down Reservoir name.
        t : str
            Time period.
            
        Returns
        -------
        poi.ConstraintIndex
            Lower bound of spillflow of reservoir.
        """
        model = self.model
        lhs = model.spillflow[r, t] - model.param['spillflow_down'][r, t]
        return model.add_linear_constraint(lhs, poi.Geq, 0)

    def import_export_up_rule(self, source: str, target: str, t: str) -> poi.ConstraintIndex:
        """Upper bound for import/export flow.

        Parameters
        ----------
        source : str
            source_stnm (外部水源/调出水水库)
        target : str
            target_reservoir_stnm (目标水库)
        t : str
            时间(YYYY-MM)

        Returns
        -------
        poi.ConstraintIndex
            Upper bound constraint for importflow[source, target, time].
        """
        model = self.model
        # 仅对已配置的三元组施加上限，集合已过滤
        up_val = model.param['import_export_up'][(source, target, t)]
        lhs = model.importflow[source, target, t] - up_val
        return model.add_linear_constraint(lhs, poi.Leq, 0)

    def import_export_down_rule(self, source: str, target: str, t: str) -> poi.ConstraintIndex:
        """Lower bound for import/export flow.

        Parameters
        ----------
        source : str
            source_stnm (外部水源/调出水水库)
        target : str
            target_reservoir_stnm (目标水库)
        t : str
            时间(YYYY-MM)

        Returns
        -------
        poi.ConstraintIndex
            Lower bound constraint for importflow[source, target, time].
        """
        model = self.model
        # 仅对已配置的三元组施加下限，集合已过滤
        down_val = model.param['import_export_down'][(source, target, t)]
        lhs = model.importflow[source, target, t] - down_val
        return model.add_linear_constraint(lhs, poi.Geq, 0)

